# 因数平方和
from leetcode import test_function as tf


def sum_square_factors(n):
    """我的方法, 暴力法, 一个用例都没有通过"""
    MOD = 1e9 + 7
    squares = [0] * (n + 1)

    def get_factors(num):
        if num == 1:
            return [1]
        factors = {1, num}
        for i in range(2, num // 2 + 1):
            if num % i == 0:
                factors.add(i)
                factors.add(num // i)
        return factors

    def sum_square(arr):
        nonlocal MOD, squares
        s = 0
        for i in arr:
            if squares[i] == 0:
                square = int(i * i % MOD)
                s += square
                squares[i] = square
            else:
                s += squares[i]
        return s

    res = 0
    for i in range(1, n + 1):
        res += sum_square(get_factors(i))

    return res


def sum_square_factors_opt(n):
    """ 通过/全部 = 2/9
    解题思路:
    反向思考: 如果a是b的因数, 那么b一定是a的倍数
    如果我们对一个数字n求其因数, 时间复杂度为 O(n), 如果1-n的每一个数字逐个求其因数, 那么时间复杂度为O(n^2), 将花费大量的时间.
    但是反过来, 我们遍历a = range(1, n/2), 并乘上一个倍数b = range(2, n/a), 那么一定可以获得n的所有因子.
    例如当n = 12:
    g(12) = f(1) + f(2) + f(3) + .... f(12)
    =1 * 12 + 2^2 * 6 + 3^2 * 4 + 4^2 * 3 + 5^2 * 2 + 6^2 * 2 + 7^2 * 1 + 8^2 * 1 + 9^2 * 1 + 10^2 + 11^2 * 1 + 12^2 * 1
    =12 + 24 + 36 + 48 + 50 + 72 + 49 + 64 + 81 + 100 + 121 + 144 = 801
    """
    MOD = 1e9 + 7
    res = 0
    for i in range(1, n + 1):
        res += (i * i) * (n // i)
        if res > MOD:
            res = res % MOD
    return int(res)


if __name__ == '__main__':
    inp = [{"n": 12}, {"n": 100000}]
    out = [801, 680584257]
    # tf(sum_square_factors, inp, out)
    tf(sum_square_factors_opt, inp, out)
    # print(sum_square_factors(12))
    # print(sum_square_factors_opt(100000))
